import streamlit as st
import torch
import matplotlib.pyplot as plt



st.set_page_config(page_title="tensor_autograd", layout="centered")
st.title("🚀 PyTorch Tensor自动求导")
st.caption("作者：何双新 ｜ 环境：Tensor高级操作+自动求导")

st.markdown("""
在本页面，你将通过动手实践，掌握以下技能：
- Tensor 广播（Broadcasting）
- 自动求导（Autograd）
""")

st.divider()

# --------------------------
# 📦 Part 1. Tensor 广播实验
# --------------------------
st.header("📦 Part 1｜Tensor 广播（Broadcasting）")

st.markdown("创建两个不同shape的Tensor，观察自动广播效果：")

a_rows = st.number_input("Tensor A - 行数", min_value=1, max_value=10, value=3)
a_cols = st.number_input("Tensor A - 列数", min_value=1, max_value=10, value=1)

b_rows = st.number_input("Tensor B - 行数", min_value=1, max_value=10, value=1)
b_cols = st.number_input("Tensor B - 列数", min_value=1, max_value=10, value=4)

if st.button("执行广播运算 ➡️"):
    A = torch.rand(int(a_rows), int(a_cols))
    B = torch.rand(int(b_rows), int(b_cols))
    try:
        C = A + B
        st.success("广播成功！结果shape:")
        st.write(f"A shape: {A.shape}, B shape: {B.shape}, C shape: {C.shape}")
        st.code(C)
    except Exception as e:
        st.error(f"广播失败，错误信息: {str(e)}")

st.divider()

# --------------------------
# 📈 Part 2. Autograd自动求导实验
# --------------------------
st.header("📈 Part 2｜Autograd 自动求导")

st.markdown("""
定义一个二次函数 `y = x² + 3x + 1` ，输入x值，自动求导出dy/dx。
""")

x_value = st.number_input("输入 x 的值", value=2.0)

if st.button("求导计算 🎯"):
    # 创建带梯度的Tensor
    x = torch.tensor(x_value, requires_grad=True)

    # 定义函数
    y = x ** 2 + 3 * x + 1

    # 反向传播
    y.backward()

    # 展示梯度
    st.success(f"y = x² + 3x + 1, 当x={x_value}时，dy/dx = {x.grad.item():.4f}")

    # 画出函数曲线和切线
    x_plot = torch.linspace(x_value - 3, x_value + 3, steps=100)
    y_plot = x_plot ** 2 + 3 * x_plot + 1

    # 切线
    slope = x.grad.item()
    tangent_line = slope * (x_plot - x_value) + y.item()

    fig, ax = plt.subplots(figsize=(8, 5))
    ax.plot(x_plot, y_plot, label='y = x² + 3x + 1')
    ax.plot(x_plot, tangent_line, '--', label=f'切线 (斜率={slope:.2f})')
    ax.scatter([x_value], [y.item()], color='red', label='切点')
    ax.legend()
    ax.set_title('函数图像与切线')
    ax.grid(True)

    st.pyplot(fig)

st.divider()
st.caption("安徽智加数字科技有限公司 · 技术学习组出品 🚀")
